'''
Author: Jackiu
Date: 2021-05-10 20:56:29
LastEditors: Jackiu
LastEditTime: 2021-05-10 21:09:06
Description: file content
'''
import json
import copy

data_root = '/media/Store2/hlp/datasets/mmw'
annotation_path = f"{data_root}/annotations"
all_data_json = {}


def print_data_info(data_json: dict, title: str):
    print(title, " dataset info -----")
    print("images: ", len(data_json['images']))
    print("annotations: ", len(data_json['annotations']))
    print("annotations / images: ", len(data_json['annotations']) / len(data_json['images']))
    print("")


def get_all_data_json():
    # load json data
    train_data_json = json.load(open(f"{annotation_path}/train.json"))
    val_data_json = json.load(open(f"{annotation_path}/val.json"))
    test_data_json = json.load(open(f"{annotation_path}/test_full.json"))

    # concat json data
    all_data_json = copy.deepcopy(train_data_json)
    all_data_json['images'].extend(val_data_json['images'])
    all_data_json['images'].extend(test_data_json['images'])
    all_data_json['annotations'].extend(val_data_json['annotations'])
    all_data_json['annotations'].extend(test_data_json['annotations'])

    # print json data info
    print_data_info(train_data_json, "train")
    print_data_info(val_data_json, "val")
    print_data_info(test_data_json, "test")
    print_data_info(all_data_json, "all")

    # sort images by name, annotations by id
    all_data_json['images'] = sorted(all_data_json['images'], key=lambda x:x['id'])
    all_data_json['annotations'] = sorted(all_data_json['images'], key=lambda x:x['id'])

    json.dump(all_data_json, open(f"{annotation_path}/all.json", 'w'))


def shuffle_seprate_data():
    empty_data_json = copy.deepcopy(all_data_json)
    empty_data_json['images'] = []
    empty_data_json['annotations'] = []

    # 取前缀
    # 列表随机打乱排序
    # 重新划分数据集
    train_data_json = copy.deepcopy(empty_data_json)
    val_data_json = copy.deepcopy(empty_data_json)
    test_data_json = copy.deepcopy(empty_data_json)


if __name__ == "__main__":
    # get_all_data_json()
    train_data_json = json.load(open(f"{annotation_path}/train.json"))
    print_data_info(train_data_json, "train")